import pyshark
from logger import log


class PcapParser:
    def __init__(self, arg):
        self.args = arg
        self.limit = 1000000
        if arg.limit is not None:
            self.limit = int(arg.limit)
        self.payloads = None
        if self.args.payload:
            self.payloads = self.args.payload.split("/")

    def parse(self, callback):
        cap = pyshark.FileCapture(self.args.file, tshark_path=self.args.tshark, display_filter=self.args.filter,
                                  decode_as=self.args.decodeas)
        limit = 0
        for packet in cap:
            if self.payloads is None or len(self.payloads) == 0:
                callback(packet[packet.transport_layer].payload.raw_value, packet[packet.transport_layer].time_relative)
            for p in self.payloads[::-1]:
                if p in packet and packet[p].has_field('payload'):
                    callback(packet[p], packet[packet.transport_layer].time_relative)
                    break
            else:
                callback(packet[packet.transport_layer].payload.raw_value, packet[packet.transport_layer].time_relative)
            limit += 1
            log.info("parse packet seq: ", packet.number)
            if limit >= self.limit:
                log.warn("parse packet reached the max limit:", self.limit)
                break
        cap.close()
        log.info(f"parse file [", self.args.file, "] finished, count=", limit)
